package com.atguigu.tree;

import com.atguigu.map.Map;

import java.util.ArrayList;
import java.util.List;

public class AVLTree<K extends Comparable<K>,V> implements Map<K,V>{

    private class Node{
        public K key;
        public V value;
        public Node right;
        public Node left;
        public int height;
        public Node(K key,V value,Node left,Node right){
            this.key=key;
            this.value=value;
            this.right=right;
            this.left=left;
            this.height=1;
        }
        public Node(K key,V value){
            this(key,value,null,null);
        }

        @Override
        public String toString() {
            return "{"+key.toString()+" : "+value.toString()+"}";
        }
    }
    private Node root;
    private int size;
    public AVLTree(){
        root=null;
        size=0;
    }
    //获得node节点的高度
    private int getHeight(Node node){
        if(node==null){
            return 0;
        }
        return node.height;
    }
    //计算平衡因子
    private int getBalanceFactor(Node node){
        return getHeight(node.left)-getHeight(node.right);
    }
    //对节点y进行右旋转，返回旋转后新的根节点x
    private Node rightRotate(Node y){
        Node x=y.left;
        Node T3=x.right;
        x.right=y;
        y.left=T3;
        //更新height
        y.height=Math.max(getHeight(y.left),getHeight(y.right))+1;
        x.height=Math.max(getHeight(x.left),getHeight(x.right))+1;
        return x;
    }
    //对节点y进行向左旋转操作，返回旋转后新的根节点x
    public Node leftRotate(Node y){
        Node x=y.right;
        Node T2=x.left;
        x.left=y;
        y.right=T2;
        //更新height
        y.height=Math.max(getHeight(y.left),getHeight(y.right))+1;
        x.height=Math.max(getHeight(x.left),getHeight(x.right))+1;
        return x;
    }
    @Override
    //向二分搜索树中添加新的元素
    public void add(K key, V value) {
        root=add(root,key,value);
    }
    //向以root为根的二分搜索树中插入元素(key,value)
    //返回插入新节点后二分搜索树的根
    private Node add(Node root,K key,V value){
        if(root==null){
            size++;
            return new Node(key,value);
        }
        if(key.compareTo(root.key)<0){
            root.left=add(root.left,key,value);
        }else if(key.compareTo(root.key)>0){
            root.right=add(root.right,key,value);
        }else{//key.compareTo(root.key)==0
            root.value=value;
        }
        //更新height
        root.height=1+Math.max(getHeight(root.left),getHeight(root.right));

        //计算平衡因子
        int balanceFactor = getBalanceFactor(root);
        //RR
        if(balanceFactor>1&&getBalanceFactor(root.left)>=0){
            return rightRotate(root);
        }
        //LL
        if(balanceFactor<-1&&getBalanceFactor(root.right)<=0){
            return leftRotate(root);
        }
        //RL
        if(balanceFactor<-1&&getBalanceFactor(root.right)>0){
            root.right=rightRotate(root.right);
            return leftRotate(root);
        }
        //LR
        if(balanceFactor>1&&getBalanceFactor(root.left)<0){
            root.left=leftRotate(root.left);
            return rightRotate(root);
        }
        return root;
    }
    //判断该二叉树是否是一颗二分搜索树
    public boolean isBST(){
        List<K>list=new ArrayList<>();
        inOrder(root,list);
        for(int i=1;i<list.size();i++){
            if(list.get(i-1).compareTo(list.get(i))>0){
                return false;
            }
        }
        return true;
    }
    //二叉树中序遍历
    private void inOrder(Node node, List<K> list) {
        if(node==null){
            return ;
        }
        inOrder(node.left,list);
        list.add(node.key);
        inOrder(node.right,list);
    }
    //判断一个数是否是平衡二叉树
    public boolean isBalanced(){
        return isBalanced(root);
    }
    //判断以Node为根的二叉树是否是一颗平衡二叉树，递归算法
    private boolean isBalanced(Node root){
        if(root==null){
            return true;
        }
        if(Math.abs(getBalanceFactor(root))<=1){
            return isBalanced(root.left)&&isBalanced(root.right);
        }else{
            return false;
        }
    }

    //返回以node为根节点的二分搜索树中，key所在的节点
    private Node getNode(Node root,K key){
        if(root==null){
            return null;
        }
        if(root.key.compareTo(key)>0){
            return getNode(root.left,key);
        }else if(root.key.compareTo(key)<0){
            return getNode(root.right,key);
        }else{
            return root;
        }
    }
    private Node minimum(Node root){
        if(root.left==null){
            return root;
        }
        return minimum(root.left);
    }
    private Node removeMin(Node root){
        if(root.left==null){
            Node rightNode=root.right;
            root.right=null;
            size--;
            return rightNode;
        }
        root.left=removeMin(root.left);
        return root;
    }
    @Override
    //从二分搜索树中删除Key节点
    public V remove(K key) {
        Node node=getNode(root,key);
        if(node!=null){
            root=remove(root,key);
            return node.value;
        }
        return null;
    }
    private Node remove(Node root,K key){
        if(root==null){
            return null;
        }
        Node retNode;
        if(key.compareTo(root.key)<0){
            root.left=remove(root.left,key);
            retNode=root;
        }else if(key.compareTo(root.key)>0){
            root.right=remove(root.right,key);
            retNode=root;
        }else{
            if(root.left==null){
                Node rightNode=root.right;
                size--;
                root.right=null;
                retNode=rightNode;
            }else if(root.right==null){
                Node leftNode=root.left;
                root.left=null;
                size--;
                retNode =leftNode;
            }else {
                Node successor = minimum(root.right);
                successor.right = remove(root.right,successor.key);
                successor.left = root.left;
                root.left = root.right = null;  //help GC
                retNode = successor;
            }
        }
        //若删除的是叶子节点，则retNode=null;
        if(retNode==null){
            return null;
        }
        //更新height
        retNode.height=1+Math.max(getHeight(retNode.left),getHeight(retNode.right));

        //计算平衡因子
        int balanceFactor = getBalanceFactor(retNode);
        //RR
        if(balanceFactor>1&&getBalanceFactor(retNode.left)>=0){
            return rightRotate(retNode);
        }
        //LL
        if(balanceFactor<-1&&getBalanceFactor(retNode.right)<=0){
            return leftRotate(retNode);
        }
        //RL
        if(balanceFactor<-1&&getBalanceFactor(retNode.right)>0){
            retNode.right=rightRotate(retNode.right);
            return leftRotate(retNode);
        }
        //LR
        if(balanceFactor>1&&getBalanceFactor(retNode.left)<0){
            retNode.left=leftRotate(retNode.left);
            return rightRotate(retNode);
        }
        return retNode;
    }

    @Override
    public boolean contains(K key) {
        return getNode(root,key)!=null;
    }

    @Override
    public V get(K key) {
        Node cur=getNode(root,key);
        return cur==null?null:cur.value;
    }

    @Override
    public void set(K key, V newValue) {
        Node cur=getNode(root,key);
        if(cur!=null){
            cur.value=newValue;
        }else{
            throw new IllegalArgumentException(key+" does't exist!");
        }
    }
    @Override
    public int getSize() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size==0;
    }
}
